package net.yegong.matrix;

import org.netlib.lapack.Ilaenv;
import org.netlib.lapack.Sgeqp3;
import org.netlib.lapack.Sorgqr;
import org.netlib.util.intW;

/**
 * @author cooper
 */
public class QRDecomposition {
	public QRDecomposition(Matrix mat) {
		this.errors = MatrixHelper.eps();
		this.mat = mat;
	}

	public void setErrorsControl(float errors) {
		this.errors = MatrixHelper.eps();
	}

	public void calc() {
		int m = mat.nRow;
		int n = mat.nCol;
		int k = Math.min(m, n);
		float[] a = mat.getColumnMajorArrayCopy();
		int[] jpvt = new int[n];
		float[] tau = new float[k];
		int nb = Ilaenv.ilaenv(1, "DGEQRF", " ", m, n, -1, -1);
		int lwork = 2 * n + (n + 1) * nb;
		float[] work = new float[lwork];
		intW info = new intW(0);

		Sgeqp3.sgeqp3(m, n, a, 0, m, jpvt, 0, tau, 0, work, 0, lwork, info);

		if (k == n) {
			float[] rarray = new float[k * n];
			for (int i = 0; i < n; ++i) {
				System.arraycopy(a, i * m, rarray, i * k, k);
			}
			r = new SquareMatrix(rarray, k, Matrix.FORTRAN);
		}
		else {
			float[] rarray = MatrixHelper.copy(a, k * n);
			r = new Matrix(rarray, k, n, Matrix.FORTRAN);
		}
		r.reserveUpperTriangle();
		for (int i = 0; i < jpvt.length; ++i) {
			--jpvt[i];
		}
		p = new PermutationMatrix(jpvt, PermutationMatrix.FROM_MAP, PermutationMatrix.COLUMNS);

		Sorgqr.sorgqr(m, k, k, a, 0, m, tau, 0, work, 0, lwork, info);

		if (n > m) {
			q = new Matrix(MatrixHelper.copy(a, m * k), m, k, Matrix.FORTRAN);
		}
		else {
			q = new Matrix(a, m, k, Matrix.FORTRAN);
		}
	}

	public Matrix getQ() {
		if (q == null)
			calc();
		return q;
	}

	public Matrix getR() {
		if (r == null)
			calc();
		return r;
	}

	public PermutationMatrix getP() {
		if (p == null)
			calc();
		return p;
	}

	public int rank() {
		if (r == null)
			calc();
		float[] diag = r.getDiag().getArray();
		int count = 0;
		float eps = mat.maxAbsolute() * errors;
		eps *= Math.max(mat.getColumnsCount(), mat.getRowsCount());
		for (int i = 0; i < diag.length; ++i) {
			if (Math.abs(diag[i]) > eps) {
				++count;
			}
		}
		return count;
	}

	protected Matrix q;
	protected Matrix r;
	protected PermutationMatrix p;
	protected float errors;
	protected Matrix mat;
}
